{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "306259b2",
   "metadata": {
    "id": "TKvYiJgnYExi"
   },
   "source": [
    "This notebook provides examples to go along with the [textbook](http://manipulation.csail.mit.edu/clutter.html).  I recommend having both windows open, side-by-side!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "481c0a20",
   "metadata": {
    "id": "A4QOaw_zYLfI",
    "lines_to_end_of_cell_marker": 2
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from IPython.display import HTML, display\n",
    "from matplotlib import pyplot as plt\n",
    "from pydrake.all import (\n",
    "    AddMultibodyPlantSceneGraph,\n",
    "    Box,\n",
    "    ConnectPlanarSceneGraphVisualizer,\n",
    "    DiagramBuilder,\n",
    "    FixedOffsetFrame,\n",
    "    JointIndex,\n",
    "    Parser,\n",
    "    PlanarJoint,\n",
    "    RandomGenerator,\n",
    "    RigidTransform,\n",
    "    RotationMatrix,\n",
    "    Simulator,\n",
    "    StartMeshcat,\n",
    "    UniformlyRandomRotationMatrix,\n",
    ")\n",
    "\n",
    "from manipulation import ConfigureParser, running_as_notebook\n",
    "from manipulation.scenarios import AddShape, ycb\n",
    "from manipulation.station import MakeHardwareStation, load_scenario"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c99c5cce",
   "metadata": {},
   "outputs": [],
   "source": [
    "meshcat = StartMeshcat()\n",
    "rng = np.random.default_rng(145)  # this is for python\n",
    "generator = RandomGenerator(rng.integers(0, 1000))  # this is for c++"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b370767",
   "metadata": {
    "id": "fGbe-rJGJlF0"
   },
   "source": [
    "# Falling things (in 2D)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "016c4e7d",
   "metadata": {
    "id": "_JIDhLfN3BSe"
   },
   "outputs": [],
   "source": [
    "def clutter_gen():\n",
    "    builder = DiagramBuilder()\n",
    "\n",
    "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step=0.001)\n",
    "\n",
    "    # Add the ground.\n",
    "    ground = AddShape(\n",
    "        plant, Box(10.0, 10.0, 10.0), \"ground\", color=[0.9, 0.9, 0.9, 1.0]\n",
    "    )\n",
    "    plant.WeldFrames(\n",
    "        plant.world_frame(),\n",
    "        plant.GetFrameByName(\"ground\", ground),\n",
    "        RigidTransform([0, 0, -5]),\n",
    "    )\n",
    "\n",
    "    # Add the bricks, each attached to the world via a planar joint.\n",
    "    parser = Parser(plant)\n",
    "    ConfigureParser(parser)\n",
    "    planar_joint_frame = plant.AddFrame(\n",
    "        FixedOffsetFrame(\n",
    "            \"planar_joint_frame\",\n",
    "            plant.world_frame(),\n",
    "            RigidTransform(RotationMatrix.MakeXRotation(np.pi / 2)),\n",
    "        )\n",
    "    )\n",
    "    for i in range(20 if running_as_notebook else 2):\n",
    "        directives = f\"\"\"\n",
    "directives:\n",
    "- add_model:\n",
    "    name: object{i}\n",
    "    file: package://manipulation/hydro/061_foam_brick.sdf\n",
    "\"\"\"\n",
    "        instance = parser.AddModelsFromString(directives, \".dmd.yaml\")[0]\n",
    "        plant.AddJoint(\n",
    "            PlanarJoint(\n",
    "                f\"joint{i}\",\n",
    "                planar_joint_frame,\n",
    "                plant.GetFrameByName(\"base_link\", instance),\n",
    "                damping=[0, 0, 0],\n",
    "            )\n",
    "        )\n",
    "\n",
    "    plant.Finalize()\n",
    "\n",
    "    vis = ConnectPlanarSceneGraphVisualizer(\n",
    "        builder,\n",
    "        scene_graph,\n",
    "        xlim=[-0.6, 0.6],\n",
    "        ylim=[-0.1, 0.5],\n",
    "        show=False,\n",
    "    )\n",
    "\n",
    "    diagram = builder.Build()\n",
    "    simulator = Simulator(diagram)\n",
    "    plant_context = plant.GetMyContextFromRoot(simulator.get_mutable_context())\n",
    "\n",
    "    z = 0.1\n",
    "    for i in range(plant.num_joints()):\n",
    "        joint = plant.get_joint(JointIndex(i))\n",
    "        if isinstance(joint, PlanarJoint):\n",
    "            joint.set_pose(\n",
    "                plant_context,\n",
    "                [rng.uniform(-0.4, 0.4), z],\n",
    "                rng.uniform(-np.pi / 2.0, np.pi / 2.0),\n",
    "            )\n",
    "            z += 0.1\n",
    "\n",
    "    vis.start_recording()\n",
    "    simulator.AdvanceTo(1.5 if running_as_notebook else 0.1)\n",
    "    vis.stop_recording()\n",
    "    ani = vis.get_recording_as_animation(repeat=False)\n",
    "    display(HTML(ani.to_jshtml()))\n",
    "\n",
    "\n",
    "clutter_gen()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6ecf1b1",
   "metadata": {
    "id": "GoxjuO5PJlF4"
   },
   "source": [
    "# Falling things (in 3D)\n",
    "\n",
    "I had to decide how to visualize the results of this one for you. The mesh and\n",
    "texture map files for the YCB objects are very large, so downloading many of\n",
    "them to your browser from an online notebook felt a bit too painful. If you've\n",
    "decided to run the notebooks from source on your local machine, then go ahead\n",
    "and open meshcat before running this test to see the live simulation. For the\n",
    "cloud notebooks, I've decided to add a camera to the scene and take a picture\n",
    "after simulating for a few seconds.  After all, that's perhaps the data that\n",
    "we're actually looking for.\n",
    "\n",
    "[Note](https://stackoverflow.com/questions/73873885/meshcat-fails-to-import-png-for-cracker-box-in-drake-planar-force-control-demo) that Drake's pip installation doesn't currently include the texture maps for the YCB objects, because they're too big to fit in the `pip` wheel 100MB limit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c02a26a",
   "metadata": {
    "id": "K2oyAnHfJlF5"
   },
   "outputs": [],
   "source": [
    "def clutter_gen():\n",
    "    scenario_data = \"\"\"\n",
    "cameras:\n",
    "    main_camera:\n",
    "        name: camera0\n",
    "        depth: True\n",
    "        X_PB:\n",
    "            base_frame: world\n",
    "            translation: [0, 0, 0.8]\n",
    "            rotation: !Rpy { deg: [180, 0, 90]}\n",
    "directives:\n",
    "- add_model:\n",
    "    name: bin\n",
    "    file: package://manipulation/hydro/bin.sdf\n",
    "- add_weld:\n",
    "    parent: world\n",
    "    child: bin::bin_base\n",
    "\"\"\"\n",
    "    for i in range(10 if running_as_notebook else 2):\n",
    "        object_num = rng.integers(0, len(ycb))\n",
    "        scenario_data += f\"\"\"\n",
    "- add_model:\n",
    "    name: thing{i}\n",
    "    file: package://manipulation/hydro/{ycb[object_num]}\n",
    "\"\"\"\n",
    "    # TODO(russt): If I didn't need to add the body name, then I could use default_free_body_pose right here and avoid the second loop.\n",
    "\n",
    "    scenario = load_scenario(data=scenario_data)\n",
    "    station = MakeHardwareStation(scenario, meshcat)\n",
    "\n",
    "    simulator = Simulator(station)\n",
    "    context = simulator.get_mutable_context()\n",
    "    plant = station.GetSubsystemByName(\"plant\")\n",
    "    plant_context = plant.GetMyContextFromRoot(context)\n",
    "\n",
    "    z = 0.2\n",
    "    for body_index in plant.GetFloatingBaseBodies():\n",
    "        tf = RigidTransform(\n",
    "            UniformlyRandomRotationMatrix(generator),\n",
    "            [rng.uniform(-0.15, 0.15), rng.uniform(-0.2, 0.2), z],\n",
    "        )\n",
    "        plant.SetFreeBodyPose(plant_context, plant.get_body(body_index), tf)\n",
    "        z += 0.1\n",
    "\n",
    "    simulator.AdvanceTo(2.0 if running_as_notebook else 0.1)\n",
    "    color_image = station.GetOutputPort(\"camera0.rgb_image\").Eval(context)\n",
    "    plt.figure()\n",
    "    plt.imshow(color_image.data)\n",
    "    plt.axis(\"off\")\n",
    "\n",
    "\n",
    "clutter_gen()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.6 64-bit",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
